import pandas as pd
import matplotlib.pyplot as plt

# 股票代码
stock_codes = ['600036', '002594', '300006', '000905', '002607']

# 每支股票的购买股数（第二种策略）
purchase_shares = {
    '600036': 20.0 * 100,  # 转换为股数
    '002594': 20.0 * 100,
    '300006': 20.0 * 100,
    '000905': 1967.0 * 100,
    '002607': 20.0 * 100
}

# CSV文件路径
base_path = './data_return/'

# 初始化一个空的DataFrame，用于存储所有股票的市值数据
all_stocks_value = pd.DataFrame()
portfolio_value = pd.DataFrame()

# 遍历每个股票代码
for code in stock_codes:
    # 构造CSV文件路径
    file_name = base_path + code + '_return.csv'

    # 读取CSV文件
    df = pd.read_csv(file_name)

    # 将date列转换为日期格式
    df['date'] = pd.to_datetime(df['date'])

    # 筛选2022年7月1日至2023年6月30日的数据
    start_date = pd.Timestamp('2022-07-01')
    end_date = pd.Timestamp('2023-06-30')
    df = df[(df['date'] >= start_date) & (df['date'] <= end_date)]

    # 获取2022年7月1日的收盘价
    initial_price = df[df['date'] == start_date]['close'].values[0]

    # 计算200万可以买多少手（手数必须为整数）
    initial_shares = int(2000000 / (initial_price * 100)) * 100  # 1手 = 100股

    # 计算每个交易日的市值
    df[f'{code}_value'] = initial_shares * df['close']

    # 第二种策略：按照指定的股数购买
    df[f'{code}_value_strategy2'] = purchase_shares[code] * df['close']

    # 将市值数据添加到总DataFrame中
    if all_stocks_value.empty:
        all_stocks_value = df[['date', f'{code}_value']]
    else:
        all_stocks_value = all_stocks_value.merge(df[['date', f'{code}_value']], on='date', how='outer')
    if portfolio_value.empty:
        portfolio_value = df[['date', f'{code}_value_strategy2']]
    else:
        portfolio_value = portfolio_value.merge(df[['date', f'{code}_value_strategy2']], on='date', how='outer')

portfolio_value['strategy2_total_value'] = portfolio_value[[col for col in portfolio_value.columns if 'strategy2' in col]].sum(axis=1)

print(portfolio_value)
# 可视化
plt.figure(figsize=(14, 8))

for code in stock_codes:
    #save all_stocks_value[['date', f'{code}_value']] to csv
    new_pd = all_stocks_value[['date', f'{code}_value']]
    #change new_pd's column f'{code}_value' to close
    new_pd.rename(columns={f'{code}_value': 'close'}, inplace=True)
    new_pd.to_csv(f'./Step6Data/{code}_value.csv')

# 设置日期为索引
all_stocks_value.set_index('date', inplace=True)
for code in stock_codes:
    plt.plot(all_stocks_value.index, all_stocks_value[f'{code}_value'], label=f'Stock {code}')

#save portfolio_value[['date', 'strategy2_total_value']] to csv
new_portfolio_df = portfolio_value[['date','strategy2_total_value']]
new_portfolio_df.rename(columns={'strategy2_total_value': 'close'}, inplace=True)
new_portfolio_df.to_csv('./Step6Data/portfolio_value.csv')

portfolio_value.set_index('date', inplace=True)
plt.plot(portfolio_value.index, portfolio_value['strategy2_total_value'], label='Strategy 2')

plt.title('Stock Value Over Time (2022-07-01 to 2023-06-30)')
plt.xlabel('Date')
plt.ylabel('Value (RMB)')
plt.legend()
plt.grid(True)
plt.show()